import numpy as np
import torchvision.ops.boxes as bops
import matplotlib.pyplot as plt
import matplotlib.patches as patches

def calculate_ious(gt_boxes, pred_boxes):
    """
    计算 gt bbox 与预测 bbox 之间的 IOU
    gt_boxes 是一个 list, 其中每个元素都是一个 tensor([x1y1x2y2])
    pred_boxes 格式与 gt_boxes 相同
    """

    box_pairs = list(zip(gt_boxes, pred_boxes))

    ious = []
    for box_pair in box_pairs:
        iou = bops.box_iou(boxes1=box_pair[0].unsqueeze(0), boxes2=box_pair[1].unsqueeze(0))
        ious.append(iou)

    return np.array(ious)

def calculate_box_iou(box1, box2):
    """
    计算两个 box 之间的 IOU
    box1 与 box2 都是 tensor([x1 y1 x2 y2])
    """
    iou = bops.box_iou(boxes1=box1.unsqueeze(0), boxes2=box2.unsqueeze(0))
    return iou

def draw_bbox(image, targets, file_path=None):
    """
    在图像上绘制 bbox
    """
    fig, a = plt.subplots(1, 1)
    fig.set_size_inches(5, 5)
    a.imshow(image)

    for box in (targets['boxes']):
        x, y, width, height = box[0], box[1], box[2]-box[0], box[3]-box[1]
        rect = patches.Rectangle(
                xy=(x, y),
                width=width,
                height=height,
                linewidth=2,
                edgecolor='r',
                facecolor='none'
            )
        a.add_patch(rect)

    if not (file_path is None):
        plt.savefig(file_path)
    else:
        plt.show()

def do_nms(gt_targets, pred_targets, threshold:float=0.5):
    """
    进行非极大值抑制
    gt_targets 与 pred_targets 都是字典的 list, 每个 dict 包含 labels, scores 和 boxes 三个键

    首先计算 pred_targets 中每个框与 gt_targets 中每个框的 iou, 剔除掉 iou < threshold 的预测
    然后计算 pred_targets 中每个框之间的 iou, 剔除其中 iou >= 0.5 的, 且 scores 值较小的那个
    """
    batch_size = len(gt_targets)
    targets: list[list[dict]] = [ [] for _ in range(batch_size) ]
    for batch in range(batch_size):
        gt: dict = gt_targets[batch]
        pred: dict = pred_targets[batch]
        for j in range(len(pred)):
            pred_label: list = pred['labels'][j]
            pred_box: list = pred['boxes'][j]
            pred_score: list = pred['scores'][j]

            for k in range(len(gt)):
                gt_box: list = gt['labels'][k]
                iou = calculate_box_iou(box1=gt_box, box2=pred_box)
                if float(iou) >= threshold:
                    targets[batch].append({
                        'label': pred_label,
                        'box': pred_box,
                        'score': pred_score
                        })
                    break
        
    for batch in range(batch_size):
        targets[batch].sort(key=lambda x: x['score'])

    del_indices: list[list[int]] = [ [] for _ in range(batch_size)]
    for batch in range(batch_size):
        for i in range(len(targets[batch])):
            for j in range(batch+1, len(targets[batch])):
                if calculate_box_iou(box1=targets[batch][i], box2=targets[batch][j]) >= threshold:
                    del_indices[batch].append(i)

    for batch in range(batch_size):
        targets[batch] = [
                targets[batch][i] for i in range(len(targets[batch])) if not i in del_indices[batch]
                ]

    result: list[dict] = []
    for batch in range(batch_size):
        result.append({
            'labels': [ target['label'] for target in targets[batch] ],
            'boxes': [ target['box'] for target in targets[batch] ],
            'scores': [ target['score'] for target in targets[batch] ],
            })
    return result


if __name__ == "__main__":
    import torchvision.transforms as transforms
    from vocdataset import PascalVOCDataset
    from torchvision.io import read_image

    data_path: str = "../data/VOC2012"
    transform = transforms.Compose([ transforms.Resize((440, 500)) ])

    voc_train = PascalVOCDataset(root=data_path, year="2012", image_set="train", transform=transform)

    _, target = next(iter(voc_train))
    image = read_image(path=target['path'])
    print(f"path = {target['path'].split('/')[-1]}")
    image = image.permute(1, 2, 0)
    draw_bbox(image=image, targets=target, file_path="./test.jpg")

    print(f"{image.shape}")
